// --------------------------------------------------------------------------------------------------------------------
// <copyright file="RGBDCameraModelAgent.cs" company="Microsoft Corporation">
// The MIT License (MIT)
// 
// Copyright (c) 2014, Microsoft Corporation
// 
// Permission is hereby granted, free of charge, to any person obtaining a copy
//  of this software and associated documentation files (the "Software"), to deal
//  in the Software without restriction, including without limitation the rights
//  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
//  copies of the Software, and to permit persons to whom the Software is
//  furnished to do so, subject to the following conditions:
// 
// The above copyright notice and this permission notice shall be included in
//  all copies or substantial portions of the Software.
// 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
//  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
//  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
//  THE SOFTWARE.
// </copyright>
// --------------------------------------------------------------------------------------------------------------------
namespace Microsoft.Robotics.Vision.Visualization
{
    using System;
    using System.Reflection;
    using System.Runtime.Serialization;
    using System.Threading.Tasks;

    using Microsoft.Robotics.Numerics;
    using Microsoft.Robotics.Runtime;
    using Microsoft.Robotics.Vision.Cameras;
    using Microsoft.Robotics.Vision.Runtime.Cameras;
    using Microsoft.Robotics.Visualization;
    
    /// <summary>
    /// This model agent class prepares the RGB frame for visualization.
    /// </summary>
    /// <typeparam name="TColorParameters">The parameters for the color camera</typeparam>
    /// <typeparam name="TDepthParameters">The parameters for the depth camera</typeparam>
    [DataContract]
    public sealed class RGBDCameraModelAgent<TColorParameters, TDepthParameters> : ModelAgent<RGBDCameraMessage<TColorParameters, TDepthParameters>, RGBDCameraModel>
        where TColorParameters : ColorParameters
        where TDepthParameters : DepthParameters
    {
        /// <summary>
        /// Width of point cloud image
        /// </summary>
        [DataMember]
        private int width;

        /// <summary>
        /// Height of point cloud image
        /// </summary>
        [DataMember]
        private int height;

        /// <summary>
        /// The background buffer
        /// </summary>
        private byte[] image;

        /// <summary>
        /// The image to display
        /// </summary>
        private byte[] displayImage;

        /// <summary>
        /// The z buffer
        /// </summary>
        private short[] zbuffer;

        /// <summary>
        /// The camera pose
        /// </summary>
        private Pose pose;

        /// <summary>
        /// The scale
        /// </summary>
        private double scale;

        /// <summary>
        /// The alignment class
        /// </summary>
        private ColorDepthAlignment alignment;

        /// <summary>
        /// Initializes a new instance of the RGBDCameraModelAgent class        
        /// </summary>
        /// <param name="name">Agent name</param>
        /// <param name="producer">The name of the agent we subscribe to.</param>
        /// <param name="maxUpdateIntervalMilliseconds">The interval on which to issue data binding updates to the UI</param>
        /// <param name="width">The width of the image that the point cloud is projected onto</param>
        /// <param name="height">The height of the image that the point cloud is projected onto</param>
        public RGBDCameraModelAgent(
            string name,
            IProducer<RGBDCameraMessage<TColorParameters, TDepthParameters>> producer,
            int maxUpdateIntervalMilliseconds,
            int width = 640,
            int height = 480)
            : base(name, producer.Name, maxUpdateIntervalMilliseconds)
        {
            this.width = 640;
            this.height = 480;
        }

        /// <summary>
        /// Called on activation
        /// </summary>
        public override void OnActivated()
        {
            base.OnActivated();
            this.Model.DisplayColor = true;
            this.Model.DisplayDepth = false;
            this.Model.PointCloud = false;
            this.Model.Blend = 0;
            this.Model.MaxValue = 10000;
            this.Model.MinValue = 1.0f;
            this.Model.LimitColor = false;
            this.Model.Yaw = 0;
            this.Model.Pitch = 0;
            this.Model.Roll = 0;
            this.image = null;
            this.zbuffer = new short[this.width * this.height];
            this.pose = new Pose(new Vector3(0, 0, 0), new Quaternion(0, 0, 0, 1));
            this.scale = 0.5;
            this.alignment = null;
        }

        /// <summary>
        /// Updates the model data with the new state.  Allows for the blending of the aligned depth to color and the aligned color to depth.
        /// Can clip the color image based on depth.  Can create a 3D point cloud.
        /// </summary>
        /// <param name="state">The new state from the upstream agent.</param>
        /// <returns>true if the model was updated, false otherwise</returns>
        protected override bool UpdateModel(RGBDCameraMessage<TColorParameters, TDepthParameters> state)
        {
            if (this.alignment == null)
            {
                this.alignment = new ColorDepthAlignment(state.AlignmentType, state.DepthCamera.Parameters, state.ColorCamera.Parameters);
            }

            // Compute the alignment between the color and depth image
            this.alignment.Align(state.DepthCamera.ImageFrame, state.ColorCamera.ImageFrame, state.DepthCamera.Parameters, state.ColorCamera.Parameters);

            System.Drawing.Size size = new System.Drawing.Size();

            if (this.Model.MaxValue < this.Model.MinValue)
            {
                this.Model.MaxValue = this.Model.MinValue + 1;
            }
            else if (this.Model.MinValue > this.Model.MaxValue)
            {
                this.Model.MinValue = this.Model.MaxValue - 1;
            }

            if (this.image != null)
            {
                Parallel.For(0, this.image.Length, (i) => this.image[i] = 0);
            }

            // Compute the scalar for the depth values to correspond to normalize
            // the depth images based on the min and max values marked by the two sliders in the UI
            float spread = this.Model.MaxValue - this.Model.MinValue;
            float scalar = 255.0f / spread;

            // Create the colored 3D point cloud
            if (this.Model.PointCloud)
            {
                float fx = (float)state.ColorCamera.Parameters.Intrinsics[0, 0];
                float fy = (float)state.ColorCamera.Parameters.Intrinsics[1, 1];
                float cx = (float)state.ColorCamera.Parameters.Intrinsics[0, 2];
                float cy = (float)state.ColorCamera.Parameters.Intrinsics[1, 2];

                if (this.image == null || this.image.Length != this.width * this.height * 3)
                {
                    this.image = new byte[this.width * this.height * 3];
                }

                Parallel.For(
                    0,
                    this.width * this.height,
                    (i) =>
                    {
                        this.zbuffer[i] = short.MaxValue;
                    });

                size = new System.Drawing.Size(this.width, this.height);

                // Compute the center
                Vector3 center = new Vector3(0, 0, 3);

                // Compute the transform that rotates the point cloud based on user input
                ////CoordinateTransform3D transform = new CoordinateTransform3D(new Pose(new Vector3(0, 0, 0), new Vector3(this.Model.Pitch, this.Model.Yaw, this.Model.Roll)));

                CoordinateTransform3D transform = new CoordinateTransform3D(new Pose(new Vector3(0, 0, 0), new Vector3(this.Model.Pitch, 0, 0)));
                CoordinateTransform3D transform2 = new CoordinateTransform3D(new Pose(new Vector3(), Quaternion.FromAxisAngle(0, 1, 0, this.Model.Yaw)));
                CoordinateTransform3D finalTransform = new CoordinateTransform3D(transform.Transform(transform2.Pose));

                this.pose = transform.Transform(this.pose);
                this.pose = transform2.Transform(this.pose);

                CoordinateTransform3D poseTransfrom = new CoordinateTransform3D(this.pose);

                // Compute the scale of the point cloud based on user input
                if (this.Model.Delta > 0)
                {
                    this.scale += 0.05;
                }
                else if (this.Model.Delta < 0)
                {
                    this.scale -= 0.05;
                }

                this.Model.Yaw = 0;
                this.Model.Pitch = 0;
                this.Model.Roll = 0;
                this.Model.Delta = 0;

                // Compute the point cloud and project it onto the image plane
                Parallel.For(
                    0,
                    state.ColorCamera.ImageFrame.ImageData.Length,
                    (i) =>
                    {
                        int x = i % state.ColorCamera.ImageFrame.Width;
                        int y = i / state.ColorCamera.ImageFrame.Width;

                        Vector3? point = this.alignment.ColorPixelTo3D(x, y);
                        if (point != null)
                        {
                            point = poseTransfrom.Transform((point.Value - center) * this.scale) + center;

                            int projectedX = (int)(cx + (point.Value.X * fx / point.Value.Z));
                            int projectedY = (int)(cy + (point.Value.Y * fy / point.Value.Z));

                            if (projectedX >= 0 && projectedX < width && projectedY >= 0 && projectedY < height && point.Value.Z > 0)
                            {
                                int depthIndex = projectedY * width + projectedX;
                                if (this.zbuffer[depthIndex] > (short)(point.Value.Z * 1000.0f))
                                {
                                    int index = (projectedY * width + projectedX) * 3;
                                    this.image[index] = state.ColorCamera.ImageFrame[x, y].Red;
                                    this.image[index + 1] = state.ColorCamera.ImageFrame[x, y].Green;
                                    this.image[index + 2] = state.ColorCamera.ImageFrame[x, y].Blue;
                                    this.zbuffer[depthIndex] = (short)(point.Value.Z * 1000.0f);
                                }
                            }
                        }
                    });
            }
            else
            {
                // Display unmodifed color and blend to mapped depth
                if (this.Model.DisplayColor)
                {
                    if (this.image == null || this.image.Length != state.ColorCamera.ImageFrame.Width * state.ColorCamera.ImageFrame.Height * 3)
                    {
                        this.image = new byte[state.ColorCamera.ImageFrame.Width * state.ColorCamera.ImageFrame.Height * 3];
                        this.displayImage = new byte[this.image.Length];
                    }

                    size = new System.Drawing.Size(state.ColorCamera.ImageFrame.Width, state.ColorCamera.ImageFrame.Height);

                    Parallel.For(
                        0,
                        state.ColorCamera.ImageFrame.ImageData.Length,
                        (i) =>
                        {
                            int x = i % state.ColorCamera.ImageFrame.Width;
                            int y = i / state.ColorCamera.ImageFrame.Width;

                            double depth = this.alignment.DepthOfRgbPixel(x, y);
                            int val = (int)((float)(depth - this.Model.MinValue) * scalar);
                            float isValid = depth == 0 ? 0.0f : 1.0f;
                            float inRange = ((this.Model.LimitColor == false) || (depth > Model.MinValue && depth < Model.MaxValue)) ? 1.0f : 0.0f;

                            int index = i * 3;

                            this.image[index] = Convert.ToByte(Math.Max(Math.Min(state.ColorCamera.ImageFrame[i].Red * (1 - this.Model.Blend) * inRange + val * (this.Model.Blend) * isValid, 255), 0));
                            this.image[index + 1] = Convert.ToByte(Math.Max(Math.Min(state.ColorCamera.ImageFrame[i].Green * (1 - this.Model.Blend) * inRange + (255.0 - val) * (this.Model.Blend) * isValid, 255), 0));
                            this.image[index + 2] = Convert.ToByte(Math.Max(Math.Min(state.ColorCamera.ImageFrame[i].Blue * (1 - this.Model.Blend) * inRange, 255), 0));
                        });
                }
                else
                {
                    // Display unmodifed depth and blend to aligned color.
                    if (this.image == null || this.image.Length != state.DepthCamera.ImageFrame.Width * state.DepthCamera.ImageFrame.Height * 3)
                    {
                        this.image = new byte[state.DepthCamera.ImageFrame.Width * state.DepthCamera.ImageFrame.Height * 3];
                    }

                    size = new System.Drawing.Size(state.DepthCamera.ImageFrame.Width, state.DepthCamera.ImageFrame.Height);

                    Parallel.For(
                        0,
                        state.DepthCamera.ImageFrame.ImageData.Length,
                        (i) =>
                        {
                            int x = i % state.DepthCamera.ImageFrame.Width;
                            int y = i / state.DepthCamera.ImageFrame.Width;

                            RgbPix color = this.alignment.ColorOfDepthPixel(x, y);

                            short depth = state.DepthCamera.ImageFrame[i];
                            int val = Math.Min(255, Math.Max(0, (int)((float)(depth - this.Model.MinValue) * scalar)));
                            int index = i * 3;
                            float isValid = depth == 0 ? 0.0f : 1.0f;
                            float inRange = ((this.Model.LimitColor == false) || (depth > Model.MinValue && depth < Model.MaxValue)) ? 1.0f : 0.0f;

                            if (depth != 0)
                            {
                                this.image[index] = Convert.ToByte(Math.Max(Math.Min(color.Red * (this.Model.Blend) * inRange + val * (1 - this.Model.Blend) * isValid, 255), 0));
                                this.image[index + 1] = Convert.ToByte(Math.Max(Math.Min(color.Green * (this.Model.Blend) * inRange + (255.0f - val) * (1 - this.Model.Blend) * isValid, 255), 0));
                                this.image[index + 2] = Convert.ToByte(Math.Max(Math.Min(color.Blue * (this.Model.Blend), 255) * inRange, 0));
                            }
                            else
                            {
                                this.image[index] = 0;
                                this.image[index + 1] = 0;
                                this.image[index + 2] = 0;
                            }
                        });
                }
            }

            this.image.CopyTo(this.displayImage, 0);
            this.Model.VideoFrame = new VideoFrame()
            {
                Data = this.displayImage,
                Size = size,
                Format = SimplePixelFormat.Rgb24,
                Stride = size.Width * 3
            };

            return true;
        }
    }
}
